{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ef189c6",
   "metadata": {},
   "source": [
    "# 测试 LambdaLift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "22393164",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入必要的TVM模块和测试工具\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relax\n",
    "from tvm.script import relax as R, tir as T, ir as I\n",
    "from tvm.relax import transform\n",
    "from tvm.ir.base import assert_structural_equal\n",
    "\n",
    "# 辅助函数：检查两个IR结构是否相等\n",
    "def _check_equal(x, y):\n",
    "    tvm.ir.assert_structural_equal(x, y)\n",
    "    tvm.ir.assert_structural_equal(y, x)\n",
    "\n",
    "    xhash = tvm.ir.structural_hash(x, map_free_vars=True)\n",
    "    yhash = tvm.ir.structural_hash(y, map_free_vars=True)\n",
    "    assert xhash == yhash\n",
    "\n",
    "\n",
    "# 辅助函数：检查IR结构是否可以正确地序列化和反序列化\n",
    "def _check_save_roundtrip(x):\n",
    "    y = tvm.ir.load_json(tvm.ir.save_json(x))\n",
    "    _check_equal(x, y)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e41675da",
   "metadata": {},
   "source": [
    "## 测试 LambdaLift 能否将局部绑定的函数提升到 IRModule 顶层"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "205b448c",
   "metadata": {},
   "source": [
    "变换前的 `IRModule`：内部函数定义在 `main` 函数内部"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ba538e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x1: R.Tensor((10, 5), \"float32\"), y1: R.Tensor((10, 5), \"float32\")\n",
    "    ) -> R.Tensor((10, 5), \"float32\"):\n",
    "        # 内部定义的函数\n",
    "        @R.function\n",
    "        def inner(\n",
    "            x2: R.Tensor((10, 5), \"float32\"), y2: R.Tensor((10, 5), \"float32\")\n",
    "        ) -> R.Tensor((10, 5), \"float32\"):\n",
    "            s: R.Tensor((10, 5), \"float32\") = R.add(x2, y2)\n",
    "            return s\n",
    "\n",
    "        gv1: R.Tensor((10, 5), \"float32\") = inner(x1, y1)\n",
    "        return gv1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b55aed20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 执行 LambdaLift 变换\n",
    "after = transform.LambdaLift()(Before)\n",
    "# 验证变换后的模块包含两个函数\n",
    "assert len(after.functions) == 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76bb1120",
   "metadata": {},
   "source": [
    "`main_inner` 是被提升的内部函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "14b51bfd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner</span>(x2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        s: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(x2, y2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> s\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_inner(x1, y1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv1\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22a1a723",
   "metadata": {},
   "source": [
    "## 测试变换不会修改输入模块"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe56a1f",
   "metadata": {},
   "source": [
    "如果输出需要新的 StructInfo，必须创建新的 relax 变量。不能更新现有 relax 变量的 struct info，因为该变量可能被其他 IRModule 使用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d3b6e7b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x: R.Tensor((2, 3), \"float32\"), y: R.Tensor((2, 3), \"float32\")\n",
    "    ) -> R.Tensor((2, 3), \"float32\"):\n",
    "        @R.function\n",
    "        def outer_func(\n",
    "            c1: R.Tensor((2, 3), \"float32\")\n",
    "        ) -> R.Callable((R.Tensor((2, 3), \"float32\"),), R.Tensor((2, 3), \"float32\")):\n",
    "            @R.function\n",
    "            def inner_func(x1: R.Tensor((2, 3), \"float32\")) -> R.Tensor((2, 3), \"float32\"):\n",
    "                s: R.Tensor((2, 3), \"float32\") = R.add(x1, c1)\n",
    "                return s\n",
    "\n",
    "            return inner_func\n",
    "\n",
    "        in_call = outer_func(x)\n",
    "        res = in_call(y)\n",
    "        return res\n",
    "\n",
    "before = Before\n",
    "# 保存原始模块的副本用于比较\n",
    "copy_of_before = tvm.ir.load_json(tvm.ir.save_json(before))\n",
    "\n",
    "# 执行LambdaLift转换\n",
    "transform.LambdaLift()(before)\n",
    "\n",
    "# 验证原始模块没有被修改\n",
    "tvm.ir.assert_structural_equal(before, copy_of_before)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0f0af27",
   "metadata": {},
   "source": [
    "## 测试闭包变换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7b58517c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 执行LambdaLift转换前的IRModule\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x: R.Tensor((2, 3), \"float32\"), y: R.Tensor((2, 3), \"float32\")\n",
    "    ) -> R.Tensor((2, 3), \"float32\"):\n",
    "        @R.function\n",
    "        def outer_func(\n",
    "            c1: R.Tensor((2, 3), \"float32\")\n",
    "        ) -> R.Callable((R.Tensor((2, 3), \"float32\"),), R.Tensor((2, 3), \"float32\")):\n",
    "            @R.function\n",
    "            def inner_func(x1: R.Tensor((2, 3), \"float32\")) -> R.Tensor((2, 3), \"float32\"):\n",
    "                # inner_func引用了外部作用域的c1变量，形成闭包\n",
    "                s: R.Tensor((2, 3), \"float32\") = R.add(x1, c1)\n",
    "                return s\n",
    "\n",
    "            return inner_func\n",
    "\n",
    "        in_call = outer_func(x)\n",
    "        res = in_call(y)\n",
    "        return res\n",
    "\n",
    "after = transform.LambdaLift()(Before)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cce4d925",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner_func</span>(x1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), c1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        s: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(x1, c1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> s\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_outer_func</span>(c1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Object:\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        inner_func: R<span style=\"color: #A2F; font-weight: bold\">.</span>Object <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>make_closure(cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_inner_func, (c1,))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> inner_func\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        in_call: R<span style=\"color: #A2F; font-weight: bold\">.</span>Object <span style=\"color: #A2F; font-weight: bold\">=</span> cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_outer_func(x)\n",
       "        res: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>invoke_pure_closure(in_call, (y,), sinfo_args<span style=\"color: #A2F; font-weight: bold\">=</span>(R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">3</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>),))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> res\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fa8ae005",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证序列化和反序列化功能\n",
    "_check_save_roundtrip(after)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5512a3a",
   "metadata": {},
   "source": [
    "## 测试递归函数提升"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b4c19fbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((2, 3), \"float32\")) -> R.Tensor:\n",
    "        @R.function\n",
    "        def while_loop(\n",
    "            i: R.Tensor((), \"int32\"), s: R.Tensor((2, 3), \"float32\")\n",
    "        ) -> R.Tensor((2, 3), \"float32\"):\n",
    "            cond: R.Tensor((), \"bool\") = R.call_pure_packed(\n",
    "                \"test.vm.less\", i, R.const(10), sinfo_args=(R.Tensor((), dtype=\"bool\"))\n",
    "            )\n",
    "            c: R.Tensor((), \"int32\") = R.const(1, dtype=\"int32\")\n",
    "            if cond:\n",
    "                new_i: R.Tensor((), \"int32\") = R.add(i, c)\n",
    "                new_s: R.Tensor((2, 3), \"float32\") = R.add(s, x)\n",
    "                # 递归调用自身\n",
    "                r: R.Tensor((2, 3), \"float32\") = while_loop(new_i, new_s)\n",
    "            else:\n",
    "                r: R.Tensor((2, 3), \"float32\") = s\n",
    "            return r\n",
    "\n",
    "        gv: R.Tensor((2, 3), \"float32\") = while_loop(R.const(0), x)\n",
    "        return gv\n",
    "\n",
    "before = Before\n",
    "# 检查递归调用的格式是否正确\n",
    "assert relax.analysis.well_formed(before)\n",
    "\n",
    "# 执行LambdaLift转换\n",
    "after = transform.LambdaLift()(before)\n",
    "# 验证转换后包含两个函数\n",
    "assert len(after.functions) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cadecd7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证序列化和反序列化功能\n",
    "_check_save_roundtrip(after)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13275b3a",
   "metadata": {},
   "source": [
    "## 测试多个顶级函数的提升"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "138b32af",
   "metadata": {},
   "source": [
    "IRModule 中的 GlobalVar 名称去重是通过附加它们被提升的函数名称来实现的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5b34af43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 转换前的IRModule：两个顶级函数都包含同名的内部函数\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def glob_func_1(\n",
    "        x1: R.Tensor((10, 5), \"float32\"), y1: R.Tensor((10, 5), \"float32\")\n",
    "    ) -> R.Tensor((10, 5), \"float32\"):\n",
    "        @R.function\n",
    "        def inner(\n",
    "            x2: R.Tensor((10, 5), \"float32\"), y2: R.Tensor((10, 5), \"float32\")\n",
    "        ) -> R.Tensor((10, 5), \"float32\"):\n",
    "            s: R.Tensor((10, 5), \"float32\") = R.add(x2, y2)\n",
    "            return s\n",
    "\n",
    "        gv1: R.Tensor((10, 5), \"float32\") = inner(x1, y1)\n",
    "        return gv1\n",
    "\n",
    "    @R.function\n",
    "    def glob_func_2(\n",
    "        x1: R.Tensor((10, 5), \"float32\"), y1: R.Tensor((10, 5), \"float32\")\n",
    "    ) -> R.Tensor((10, 5), \"float32\"):\n",
    "        @R.function\n",
    "        def inner(\n",
    "            x2: R.Tensor((10, 5), \"float32\"), y2: R.Tensor((10, 5), \"float32\")\n",
    "        ) -> R.Tensor((10, 5), \"float32\"):\n",
    "            s: R.Tensor((10, 5), \"float32\") = R.add(x2, y2)\n",
    "            return s\n",
    "\n",
    "        gv1: R.Tensor((10, 5), \"float32\") = inner(x1, y1)\n",
    "        return gv1\n",
    "\n",
    "before = Before\n",
    "# 执行LambdaLift转换\n",
    "after = transform.LambdaLift()(before)\n",
    "# 验证转换后包含4个函数\n",
    "assert len(after.functions) == 4\n",
    "# 验证序列化和反序列化功能\n",
    "_check_save_roundtrip(after)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c86dfb1c",
   "metadata": {},
   "source": [
    "## 测试无局部函数的情况"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "46d74a08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 没有局部函数的IRModule\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @T.prim_func\n",
    "    def sub(\n",
    "        A: T.Buffer((16, 16), \"float32\"),\n",
    "        B: T.Buffer((16, 16), \"float32\"),\n",
    "        C: T.Buffer((16, 16), \"float32\"),\n",
    "    ) -> None:\n",
    "        for i, j in T.grid(16, 16):\n",
    "            with T.block(\"sub\"):\n",
    "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "                C[vi, vj] = A[vi, vj] - B[vi, vj]\n",
    "\n",
    "    @R.function\n",
    "    def before(c0: R.Tensor((16, 16), \"float32\"), x: R.Tensor(dtype=\"float32\", ndim=2)):\n",
    "        s = R.call_tir(Before.sub, (c0, x), R.Tensor((16, 16), dtype=\"float32\"))\n",
    "        return s\n",
    "\n",
    "before = Before\n",
    "# 执行LambdaLift转换\n",
    "after = transform.LambdaLift()(before)\n",
    "# 验证没有局部函数被提升，模块保持不变\n",
    "assert_structural_equal(after, before, map_free_vars=True)\n",
    "# 验证序列化和反序列化功能\n",
    "_check_save_roundtrip(after)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6206e7a",
   "metadata": {},
   "source": [
    "## 测试非纯函数的提升"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7660cfb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 转换前的IRModule：包含非纯内部函数\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function(pure=False)\n",
    "    def main(x: R.Tensor((), \"int32\")) -> R.Tensor((), \"int32\"):\n",
    "        # 非纯内部函数，使用R.print产生副作用\n",
    "        @R.function(pure=False)\n",
    "        def inner() -> R.Tuple:\n",
    "            y = R.print(format=\"Wow!\")\n",
    "            return y\n",
    "\n",
    "        gv1 = inner()\n",
    "        return x\n",
    "\n",
    "before = Before\n",
    "# 执行LambdaLift转换\n",
    "after = transform.LambdaLift()(before)\n",
    "# 验证转换后包含两个函数\n",
    "assert len(after.functions) == 2\n",
    "# 验证序列化和反序列化功能\n",
    "_check_save_roundtrip(after)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "61a7a946",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(pure<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>, private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner</span>() <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tuple:\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>print(format<span style=\"color: #A2F; font-weight: bold\">=</span>R<span style=\"color: #A2F; font-weight: bold\">.</span>str(<span style=\"color: #BA2121\">&quot;Wow!&quot;</span>))\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>tuple()\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(pure<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">False</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_inner()\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> x\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0edd908c",
   "metadata": {},
   "source": [
    "## 测试与全局函数同名的 lambda 函数提升"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c309cf85",
   "metadata": {},
   "source": [
    "测试提升的 lambda 名称可能不会与先前的名称冲突。模块已有名为`main_inner`的函数，该名称与LambdaLift为提升函数选择的第一个名称相同。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1b6db14a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner</span>() <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tuple:\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>tuple()\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner_0</span>(x2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        s: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(x2, y2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> s\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        gv1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_inner_0(x1, y1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv1\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 转换前的IRModule：已包含一个名为main_inner的全局函数\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x1: R.Tensor((10, 5), \"float32\"), y1: R.Tensor((10, 5), \"float32\")\n",
    "    ) -> R.Tensor((10, 5), \"float32\"):\n",
    "        @R.function\n",
    "        def inner(\n",
    "            x2: R.Tensor((10, 5), \"float32\"), y2: R.Tensor((10, 5), \"float32\")\n",
    "        ) -> R.Tensor((10, 5), \"float32\"):\n",
    "            s: R.Tensor((10, 5), \"float32\") = R.add(x2, y2)\n",
    "            return s\n",
    "\n",
    "        gv1: R.Tensor((10, 5), \"float32\") = inner(x1, y1)\n",
    "        return gv1\n",
    "\n",
    "    # 已存在的全局函数，名称为main_inner\n",
    "    @R.function\n",
    "    def main_inner():\n",
    "        return R.tuple()\n",
    "\n",
    "# 执行LambdaLift转换\n",
    "after = transform.LambdaLift()(Before)\n",
    "after.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02ffa293",
   "metadata": {},
   "source": [
    "## 测试由内部函数定义的符号变量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "28db05a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner</span>(x2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        m <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        sum_inner: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((n, m), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(x2, y2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> sum_inner\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        sum_main: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_inner(x1, y1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> sum_main\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 转换前的IRModule：内部函数使用符号变量定义张量形状\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x1: R.Tensor((10, 5), \"float32\"), y1: R.Tensor((10, 5), \"float32\")\n",
    "    ) -> R.Tensor((10, 5), \"float32\"):\n",
    "        @R.function\n",
    "        def inner(x2: R.Tensor((\"n\", \"m\"), \"float32\"), y2: R.Tensor((\"n\", \"m\"), \"float32\")):\n",
    "            sum_inner = R.add(x2, y2)\n",
    "            return sum_inner\n",
    "\n",
    "        sum_main = inner(x1, y1)\n",
    "        return sum_main\n",
    "\n",
    "# 执行LambdaLift转换\n",
    "After = transform.LambdaLift()(Before)\n",
    "After.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bd9405e",
   "metadata": {},
   "source": [
    "## 测试由外部函数定义的符号变量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8e44ab50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function(private<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main_inner</span>(x2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y2: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        m <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        sum_inner: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((n, m), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>add(x2, y2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> sum_inner\n",
       "\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), y1: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #BA2121\">&quot;n&quot;</span>, <span style=\"color: #BA2121\">&quot;m&quot;</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        n <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        m <span style=\"color: #A2F; font-weight: bold\">=</span> T<span style=\"color: #A2F; font-weight: bold\">.</span>int64()\n",
       "        cls <span style=\"color: #A2F; font-weight: bold\">=</span> Module\n",
       "        sum_main: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((n, m), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> cls<span style=\"color: #A2F; font-weight: bold\">.</span>main_inner(x1, y1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> sum_main\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 转换前的IRModule：内部函数使用外部函数定义的符号变量\n",
    "@I.ir_module\n",
    "class Before:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x1: R.Tensor((\"n\", \"m\"), \"float32\"), y1: R.Tensor((\"n\", \"m\"), \"float32\")\n",
    "    ) -> R.Tensor((\"n\", \"m\"), \"float32\"):\n",
    "        # 在外部函数中定义符号变量\n",
    "        n = T.int64()\n",
    "        m = T.int64()\n",
    "\n",
    "        @R.function\n",
    "        def inner(x2: R.Tensor((n, m), \"float32\"), y2: R.Tensor((n, m), \"float32\")):\n",
    "            sum_inner = R.add(x2, y2)\n",
    "            return sum_inner\n",
    "\n",
    "        sum_main = inner(x1, y1)\n",
    "        return sum_main\n",
    "\n",
    "# 预期的IRModule结构：提升后的函数正确处理外部定义的符号变量\n",
    "# 执行LambdaLift转换\n",
    "After = transform.LambdaLift()(Before)\n",
    "After.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5a501f7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
